import numpy as np
import matplotlib.pyplot as plt
import time


class NBClassifier(object):
    def __init__(self):
        self.class_prob_mat = None
        self.condition_prob_mat = None
        self.wordSet = None

    def train(self, data_set, onehot_labels: np.ndarray):
        data_array = np.array(data_set, np.float32)
        assert len(data_array.shape) == 2, 'data_set必须是维数为2的向量'
        labels_array = np.array(onehot_labels)
        assert len(labels_array.shape) == 1
        assert data_array.shape[0] == labels_array.shape[0], '样本和标注的长度必须相等'

        train_size = data_array.shape[0]
        num_feature = data_array.shape[1]

        self.class_prob_mat = np.log(labels_array.sum(axis=0).astype(np.float32) / train_size)
        self.condition_prob_mat = np.zeros((labels_array.shape[1], num_feature), np.float32)
        for i in range(self.condition_prob_mat.shape[0]):
            class_i_data = [data_array[x] for x in range(train_size) if labels_array[x][i] == 1]
            class_i_data = np.array(class_i_data)
            class_i_prob = class_i_data.sum(axis=0) + np.ones(self.condition_prob_mat.shape[1])
            class_i_prob = np.log(class_i_prob.astype(np.float32) / class_i_prob.sum())
            self.condition_prob_mat[i] = class_i_prob

    def predict(self, input_vec):
        prob = np.zeros_like(self.class_prob_mat)
        fun = lambda x: np.sum(input_vec * self.condition_prob_mat[x]) + self.class_prob_mat[x]
        for x in range(prob.shape[0]):
            prob[x] = fun(x)

        predict_class = np.argmax(prob)
        return predict_class

    def create_word_set(self, text_lines):
        wordSet = set()
        for line in text_lines:
            for word in line:
                wordSet.add(word)
        self.wordSet = list[wordSet]

    def word_to_vec(self, input_text):
        ret = [0] * len(self.wordSet)
        for word in input_text:
            if word in self.wordSet:
                ret[self.wordSet.index(word)] += 1
        return ret

    @staticmethod
    def text_parse(big_string):
        import re
        list_of_tokens = re.split(r'\W*', big_string)
        return [tok.lower() for tok in list_of_tokens if len(tok) > 2]
